package com.kujie.ai.utils;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.github.guoshiqiufeng.dify.chat.dto.response.ChatMessageSendCompletionResponse;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * 流式响应工具类
 *
 * @author check
 * @date 2025-08-06
 */
public class StreamUtils {
    private static final ObjectMapper objectMapper = new ObjectMapper();

    public static Flux<String> parseWorkflowOutput(Flux<ChatMessageSendCompletionResponse> workflowSseFlux) {
        return workflowSseFlux
                .flatMap(response -> {
                    // 检查事件类型是否为 node_finished
                    if ("node_finished".equals(response.getEvent())) {
                        // 获取 data 字段
                        Object dataObj = response.getData();
                        if (dataObj != null) {
                            // 尝试将 data 转换为 Map
                            Map<String, Object> data = convertToMap(dataObj);
                            if (data != null) {
                                StringBuilder results = new StringBuilder();
                                if (data.get("nodeType").equals("llm")) {
                                    // 获取 outputs 字段
                                    Object outputsObj = data.get("outputs");
                                    if (outputsObj instanceof Map) {
                                        Map<String, Object> outputs = (Map<String, Object>) outputsObj;


                                        // 提取 answer 字段
                                        Object text = outputs.get("text");
                                        if (text != null) {
                                            String textString = text.toString();
                                            // 清理 answer 内容（移除<think>标签等）
                                            String cleanedAnswer = cleanAnswer(textString);
                                            if (!cleanedAnswer.isEmpty()) {
                                                results.append(cleanedAnswer);
                                            }
                                        }

                                    }
                                } else if (data.get("nodeType").equals("tool")){
                                    // 获取 outputs 字段
                                    Object outputsObj = data.get("outputs");
                                    if (outputsObj instanceof Map) {
                                        Map<String, Object> outputs = (Map<String, Object>) outputsObj;
                                        // 提取 msg 字段
                                        Object json = outputs.get("json");
                                        if (json instanceof ArrayList<?> jsonArray) {
                                            Map<String,Object> jsonMap = (Map<String,Object> )jsonArray.get(0);
                                            Object msg = jsonMap.get("msg");
                                            if (msg != null) {
                                                results.append(msg.toString());
                                            }
                                        }

                                    }
                                }

                                // 将结果按每5个字符切分
                                if (!results.isEmpty()) {
                                    List<String> chunks = splitIntoChunks(results.toString());
                                    return Flux.fromIterable(chunks);
                                }
                            }
                        }
                    }
                    // 对于非 node_finished 事件，返回空 Flux
                    return Flux.empty();
                });
    }

    // 将字符串按指定长度切分成 List
    private static List<String> splitIntoChunks(String text) {
        List<String> chunks = new ArrayList<>();
        for (int i = 0; i < text.length(); i += 5) {
            int endIndex = Math.min(i + 5, text.length());
            chunks.add(text.substring(i, endIndex));
        }
        return chunks;
    }

    // 将对象转换为 Map
    private static Map<String, Object> convertToMap(Object obj) {
        if (obj instanceof Map) {
            return (Map<String, Object>) obj;
        }

        try {
            // 尝试使用 Jackson 将对象转换为 Map
            return objectMapper.convertValue(obj, Map.class);
        } catch (Exception e) {
            // 转换失败，返回 null
            return null;
        }
    }

    public static String cleanAnswer(String answer) {
        if (answer == null || answer.isEmpty()) {
            return answer;
        }

        // 移除<think>标签
        String cleaned = answer.replaceAll("<think>", "\n\n").trim();

        // 移除<</think>>标签
        cleaned = cleaned.replaceAll("</think>", "\n\n").trim();

        // 移除多余的空行
        cleaned = cleaned.replaceAll("\n\\s*\n", "\n\n").trim();

        return cleaned;
    }

    // 解析 output 字段的方法
    private static String parseOutput(String output) {
        if (output == null || output.isEmpty()) {
            return "";
        }

        try {
            // 将字符串转换为 JSON 节点
            JsonNode jsonArray = objectMapper.readTree(output);

            // 遍历数组中的每个元素
            StringBuilder result = new StringBuilder();
            if (jsonArray.isArray()) {
                for (JsonNode item : jsonArray) {
                    JsonNode msgNode = item.get("msg");
                    if (msgNode != null) {
                        if (!result.isEmpty()) {
                            result.append("\n");
                        }
                        result.append(msgNode.asText());
                    }
                }
            } else if (jsonArray.isObject()) {
                // 如果是单个对象而不是数组
                JsonNode msgNode = jsonArray.get("msg");
                if (msgNode != null) {
                    result.append(msgNode.asText());
                }
            }

            return result.toString();
        } catch (Exception e) {
            // 如果解析失败，返回原始字符串
            return output;
        }
    }
}
